import os
from typing import Dict, List, Tuple

import pandas as pd

from iflytech_assistant.spark_client import instruct

from .dataclasses import FiveSuggestions, PreferenceTag, RefineWord
from .markdown_prompt_parser import PromptTemplate

__all__ = ["generate", "generate_response"]

templates_dir = os.path.join(os.path.dirname(__file__), "templates")
template_file = os.path.join(templates_dir, "review_guru.md")
prompt_template = PromptTemplate.from_markdown_text(open(template_file).read())

SYSTEM_PROMPT = prompt_template.system_prompt_template.raw_str
POLISH_USER_PROMPT = prompt_template.polish_user_prompt_template.raw_str


def build_prompt(
    user_input: str,
    target: str,
    tags: List[PreferenceTag],
    mode: str,
    refine_words: RefineWord = None,
) -> Tuple[str, str]:
    style_prompt = "\n".join([f"- {tag.tag}: {tag.style}" for tag in tags if tag.style])
    target_prompt = "\n".join(
        [f"- {tag.tag}: {tag.target}" for tag in tags if tag.target]
    )

    if refine_words:
        style_prompt = refine_words.word

    system_prompt = SYSTEM_PROMPT.format(target_block=target_prompt, style=style_prompt)

    user_prompt = POLISH_USER_PROMPT.format(
        target=target,
        user_input=user_input,
    )
    return system_prompt, user_prompt


def generate(
    user_input,
    target: str,
    tags: List[PreferenceTag],
    mode: str,
    refine_words: RefineWord = None,
):
    system_prompt, user_prompt = build_prompt(
        user_input, target, tags, mode, refine_words
    )
    llm_response = instruct(user_prompt, system_prompt, prefix="", streaming=False)

    suggestions = FiveSuggestions.from_llm_response(llm_response)

    return system_prompt, user_prompt, llm_response, suggestions


def generate_response(
    user_input: str,
    target: str,
    tags: List[PreferenceTag],
    refine_words: RefineWord,
    mode: str = "polish",
) -> Tuple[FiveSuggestions, FiveSuggestions]:
    system_prompt, user_prompt, _, suggestions = generate(
        user_input, target, tags, mode
    )

    refine_prompt, _, _, refine_suggestions = generate(
        user_input, target, tags, mode, refine_words
    )
    return (
        suggestions,
        refine_suggestions,
        system_prompt,
        user_prompt,
        refine_prompt,
    )


def load_tags_from_excel(
    file_path: str,
) -> Tuple[Dict[str, List[PreferenceTag]], Dict[str, List[RefineWord]]]:
    # read all sheet names
    xls = pd.ExcelFile(file_path)
    sheets = xls.sheet_names

    assert "targets" in sheets, "targets sheet not found"
    assert "refine_words" in sheets, "refine_words sheet not found"

    refine_words_sheet = xls.parse("refine_words")
    # remove first row
    refine_words_sheet = refine_words_sheet.iloc[1:]
    # column names 高情商-{target} -> target
    refine_words_sheet.columns = refine_words_sheet.columns.str.extract(
        r"评价达人-(.*)"
    )[0]

    refine_words = {}
    for idx, col in enumerate(refine_words_sheet.columns):
        if not isinstance(col, str) or "Unnamed" in col:
            continue
        words = refine_words_sheet[col].dropna().tolist()
        types = refine_words_sheet.iloc[:, idx + 1].dropna().tolist()
        refine_words[col] = [
            RefineWord(word=word, type=types[i]) for i, word in enumerate(words)
        ]

    target_sheet = xls.parse("targets")
    target_objs = {}
    # skip first row

    for idx, row in target_sheet.iterrows():
        if idx == 0:
            continue
        target_objs[row.iloc[0]] = [
            PreferenceTag(
                tag=row.iloc[0],
                target=row.iloc[1],
            )
        ]

    return target_objs, refine_words
